import pandas as pd
import numpy as np
import backtrader as bt
import backtrader.analyzers as btanalyzers
import matplotlib.pyplot as plt


# 1. 数据预处理：前复权计算
def adjust_price(df):
    """
    计算前复权价格
    公式：前复权价格 = 当日价格 × 当日复权因子 / 最新复权因子
    """
    # 计算复权因子（这里简化处理，实际应从tushare获取复权因子）
    df['adj_factor'] = df['close'] / df['pre_close'].shift(-1).fillna(1)
    df['adj_factor'] = df['adj_factor'].cumprod().shift(1).fillna(1)

    # 计算前复权价格
    for col in ['open', 'high', 'low', 'close']:
        df[col] = df[col] * df['adj_factor']

    return df


# 2. 网格策略实现
class GridStrategy(bt.Strategy):
    params = (
        ('initial_invest', 100000),
        ('grid_step', 0.05),
        ('grid_amount', 10000),
        ('slippage_perc', 0.01),
    )

    def __init__(self):
        super().__init__()
        self.base_price = None
        self.last_level = 0
        self.order = None

    def log(self, txt, dt=None, is_warning=False):
        dt = dt or self.datas[0].datetime.date(0)
        prefix = "⚠️ 警告：" if is_warning else ""
        print(f'{dt.isoformat()}, {prefix}{txt}')

    def check_risk(self, size, is_buy):
        """风控检查并返回可交易数量"""
        if is_buy:
            available_cash = self.broker.getcash()
            needed_cash = abs(size) * self.data.close[0] * (1 + self.p.slippage_perc)

            if available_cash < needed_cash:
                self.log(f"现金不足，需要{needed_cash:.2f}元，可用{available_cash:.2f}元",
                         is_warning=True)
                return 0  # 返回0表示无法交易
            return size

        else:  # 卖出检查
            position = self.getposition()
            if position.size < abs(size):
                self.log(f"持仓不足，需要卖出{abs(size):.2f}股，当前持有{position.size:.2f}股",
                         is_warning=True)
                return 0
            return size

    def next(self):
        if len(self) == 1:
            return

        # 初始建仓
        if self.base_price is None:
            self.base_price = self.data.close[0]
            invest_amount = min(self.p.initial_invest, self.broker.getcash())
            size = invest_amount / self.data.close[0]
            self.order = self.buy(size=size)
            self.log(f"初始建仓 {size:.2f}股 @ {self.data.close[0]:.2f}元")
            return

        # 网格交易逻辑
        price_change = (self.data.close[0] - self.base_price) / self.base_price
        current_level = int(price_change // self.p.grid_step)
        level_diff = current_level - self.last_level

        if level_diff == 0:
            return

        # 计算理论交易量
        price = self.data.close[0] * (1 - np.sign(level_diff) * self.p.slippage_perc)
        target_size = abs(level_diff) * self.p.grid_amount / price

        # 执行风控检查
        adjusted_size = self.check_risk(
            size=target_size * np.sign(level_diff),
            is_buy=(level_diff < 0)
        )

        # 执行实际交易
        if adjusted_size != 0:
            if level_diff > 0:
                self.order = self.sell(size=adjusted_size)
                self.log(f"卖出 {adjusted_size:.2f}股 @ {price:.2f}元")
            else:
                self.order = self.buy(size=adjusted_size)
                self.log(f"买入 {adjusted_size:.2f}股 @ {price:.2f}元")
            self.last_level = current_level

    def notify_order(self, order):
        if order.status in [order.Completed]:
            cash = self.broker.getcash()
            position = self.getposition()
            self.log(f"交易完成 | 持仓:{position.size:.2f}股, 现金:{cash:.2f}元")
# 3. 主回测函数
def run_backtest():
    # 创建回测引擎
    cerebro = bt.Cerebro()

    # 设置初始资金
    cerebro.broker.setcash(200000)

    # 设置滑点
    cerebro.broker.set_slippage_perc(perc=0.001)

    # 无手续费
    cerebro.broker.setcommission(commission=0.0)

    # 加载数据
    df = pd.read_csv('../data/513180.SH_daily_data.csv')

    # 前复权处理
    df = adjust_price(df)

    # 转换为Backtrader数据格式
    df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
    df.set_index('trade_date', inplace=True)
    df.sort_index(inplace=True)

    data = bt.feeds.PandasData(
        dataname=df,
        datetime=None,  # 索引列是日期
        open=2,  # open在列2 (0-indexed)
        high=3,
        low=4,
        close=5,
        volume=8,
        openinterest=-1
    )
    cerebro.adddata(data)

    # 添加策略
    cerebro.addstrategy(GridStrategy)

    # 添加分析器
    cerebro.addanalyzer(btanalyzers.SharpeRatio, _name='sharpe')
    cerebro.addanalyzer(btanalyzers.DrawDown, _name='drawdown')
    cerebro.addanalyzer(btanalyzers.Returns, _name='returns')

    # 运行回测
    print('初始资金: %.2f' % cerebro.broker.getvalue())
    results = cerebro.run()
    print('最终资金: %.2f' % cerebro.broker.getvalue())

    # 打印分析结果
    strat = results[0]
    print('夏普比率:', strat.analyzers.sharpe.get_analysis()['sharperatio'])
    print('最大回撤: %.2f%%' % strat.analyzers.drawdown.get_analysis()['max']['drawdown'])
    print('年化收益率: %.2f%%' % (strat.analyzers.returns.get_analysis()['rnorm100']))

    # 绘制结果
    cerebro.plot(style='candle', volume=False)


if __name__ == '__main__':
    run_backtest()